#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RAG系统依赖安装脚本
检查和安装必需的Python包
"""

import subprocess
import sys


def install_package(package_name, version=None):
    """安装Python包"""
    try:
        if version:
            package_spec = f"{package_name}>={version}"
        else:
            package_spec = package_name
        
        print(f"正在安装 {package_spec}...")
        
        result = subprocess.run([
            sys.executable, "-m", "pip", "install", package_spec
        ], capture_output=True, text=True, check=True)
        
        print(f"✅ {package_name} 安装成功")
        return True
        
    except subprocess.CalledProcessError as e:
        print(f"❌ {package_name} 安装失败:")
        print(f"   错误输出: {e.stderr}")
        return False
    except Exception as e:
        print(f"❌ 安装 {package_name} 时出现异常: {e}")
        return False


def check_package(package_name):
    """检查包是否已安装"""
    try:
        __import__(package_name)
        print(f"✅ {package_name} 已安装")
        return True
    except ImportError:
        print(f"❌ {package_name} 未安装")
        return False


def main():
    """主函数"""
    print("🔧 RAG系统依赖检查和安装")
    print("=" * 50)
    
    # 必需的依赖包
    dependencies = [
        ("requests", "2.31.0"),
        ("pymilvus", "2.3.0"),
        ("flask", "2.3.0"),
        ("torch", "1.12.1"),
        ("sentence-transformers", "2.2.2"),
        ("numpy", "1.21.0")
    ]
    
    print("📋 检查依赖包...")
    missing_packages = []
    
    for package_name, version in dependencies:
        if not check_package(package_name.replace("-", "_")):  # 有些包名不同
            missing_packages.append((package_name, version))
    
    if not missing_packages:
        print("\n🎉 所有依赖包都已安装！")
        return
    
    print(f"\n📦 需要安装 {len(missing_packages)} 个包:")
    for package_name, version in missing_packages:
        print(f"   - {package_name} >= {version}")
    
    # 询问是否安装
    try:
        choice = input(f"\n是否现在安装这些包？(y/n): ").lower().strip()
        
        if choice == 'y':
            print("\n🚀 开始安装依赖包...")
            
            success_count = 0
            for package_name, version in missing_packages:
                if install_package(package_name, version):
                    success_count += 1
            
            print(f"\n📊 安装结果: {success_count}/{len(missing_packages)} 个包安装成功")
            
            if success_count == len(missing_packages):
                print("🎉 所有依赖包安装完成！")
                print("\n现在您可以运行:")
                print("   python3 rag_demo_fixed.py")
                print("   python3 rag_chatbot.py")
                print("   python3 rag_web_chat.py")
            else:
                print("⚠️ 部分包安装失败，请手动安装剩余的包")
        else:
            print("\n📝 手动安装命令:")
            print("pip install", end="")
            for package_name, version in missing_packages:
                print(f" {package_name}>={version}", end="")
            print()
    
    except KeyboardInterrupt:
        print("\n⏹️ 安装被中断")
    except Exception as e:
        print(f"\n❌ 安装过程中出现错误: {e}")


if __name__ == "__main__":
    main()